"""
Multi-field PDE framework for coupled systems

This module provides a generic framework for defining and solving coupled PDE systems
with multiple interacting fields.
"""

using ..Fields
using ..Grids
using LinearAlgebra

"""
    MultiFieldPDE <: AbstractPDE

Generic framework for coupled PDE systems with multiple fields.

# Examples
```julia
# Reaction-diffusion system: ∂u/∂t = D₁∇²u + f(u,v), ∂v/∂t = D₂∇²v + g(u,v)
function reaction_diffusion(fields, t, params)
    u, v = fields[1], fields[2]
    D1, D2 = params[:D1], params[:D2]
    
    # Diffusion terms
    dudt = D1 * laplacian(u)
    dvdt = D2 * laplacian(v) 
    
    # Reaction terms
    u_data, v_data = data(u), data(v)
    dudt += ScalarField(u_data .* (1 .- u_data .- 2*v_data), grid(u))
    dvdt += ScalarField(v_data .* (u_data .- v_data), grid(v))
    
    return FieldCollection([dudt, dvdt])
end

pde = MultiFieldPDE(reaction_diffusion, Dict(:D1 => 0.1, :D2 => 0.05))
```
"""
struct MultiFieldPDE <: AbstractPDE
    system_function::Function    # (fields, t, params) -> field_derivatives
    parameters::Dict{Symbol, Any}
    field_names::Vector{String}
    bc_conditions::Dict{String, String}
    
    function MultiFieldPDE(system_func::Function, params::Dict{Symbol, Any}=Dict{Symbol, Any}();
                          field_names::Vector{String}=String[], 
                          bc_conditions::Dict{String, String}=Dict{String, String}())
        new(system_func, params, field_names, bc_conditions)
    end
end

"""
    evolution_rate(pde::MultiFieldPDE, state::FieldCollection, t::Float64)

Compute evolution rates for the multi-field system.
"""
function evolution_rate(pde::MultiFieldPDE, state::FieldCollection, t::Float64)
    return pde.system_function(state, t, pde.parameters)
end

"""
    ReactionDiffusionPDE <: AbstractPDE

Specialized reaction-diffusion system for two species:
    ∂u/∂t = D₁∇²u + R₁(u,v)
    ∂v/∂t = D₂∇²v + R₂(u,v)
"""
struct ReactionDiffusionPDE <: AbstractPDE
    D1::Float64              # Diffusion coefficient for species 1
    D2::Float64              # Diffusion coefficient for species 2
    reaction_function::Function  # (u, v) -> (R₁, R₂)
    bc::String
    
    function ReactionDiffusionPDE(D1, D2, reaction_func; bc="auto_periodic_neumann")
        new(D1, D2, reaction_func, bc)
    end
end

"""
    evolution_rate(pde::ReactionDiffusionPDE, state::FieldCollection, t::Float64)

Evolution rate for reaction-diffusion system.
"""
function evolution_rate(pde::ReactionDiffusionPDE, state::FieldCollection, t::Float64)
    if length(state) != 2
        throw(ArgumentError("ReactionDiffusionPDE requires exactly 2 fields: [u, v]"))
    end
    
    u, v = state[1], state[2]
    
    # Diffusion terms
    u_diff = pde.D1 * laplacian(u)
    v_diff = pde.D2 * laplacian(v)
    
    # Reaction terms
    u_data, v_data = data(u), data(v)
    R1, R2 = pde.reaction_function(u_data, v_data)
    
    u_react = ScalarField(R1, grid(u))
    v_react = ScalarField(R2, grid(v))
    
    dudt = u_diff + u_react
    dvdt = v_diff + v_react
    
    return FieldCollection([dudt, dvdt])
end

"""
    FitzHughNagumoReaction(a, b, c, I)

FitzHugh-Nagumo reaction terms for excitable media.
Returns functions for: ∂u/∂t = u(u-a)(1-u) - v + I, ∂v/∂t = b(u - cv)
"""
function FitzHughNagumoReaction(a=0.1, b=0.01, c=0.1, I=0.0)
    return (u, v) -> begin
        R1 = u .* (u .- a) .* (1.0 .- u) .- v .+ I
        R2 = b .* (u .- c .* v)
        return (R1, R2)
    end
end

"""
    SchnakenbergReaction(a, b)

Schnakenberg reaction terms for pattern formation.
Returns functions for: ∂u/∂t = a - u + u²v, ∂v/∂t = b - u²v
"""
function SchnakenbergReaction(a=0.1, b=0.9)
    return (u, v) -> begin
        u2v = u.^2 .* v
        R1 = a .- u .+ u2v
        R2 = b .- u2v
        return (R1, R2)
    end
end

"""
    GrayScottReaction(feed, kill)

Gray-Scott reaction terms for spot patterns.
Returns functions for: ∂u/∂t = -uv² + feed(1-u), ∂v/∂t = uv² - (kill+feed)v
"""
function GrayScottReaction(feed=0.04, kill=0.06)
    return (u, v) -> begin
        uv2 = u .* v.^2
        R1 = -uv2 .+ feed .* (1.0 .- u)
        R2 = uv2 .- (kill + feed) .* v
        return (R1, R2)
    end
end

"""
    BrusselatorReaction(A, B)

Brusselator reaction terms for oscillatory patterns.
Returns functions for: ∂u/∂t = A + u²v - (B+1)u, ∂v/∂t = Bu - u²v
"""
function BrusselatorReaction(A=1.0, B=3.0)
    return (u, v) -> begin
        u2v = u.^2 .* v
        R1 = A .+ u2v .- (B + 1.0) .* u
        R2 = B .* u .- u2v
        return (R1, R2)
    end
end

"""
    CoupledOscillatorPDE <: AbstractPDE

System of coupled oscillators with spatial coupling:
    ∂u/∂t = v + D₁∇²u
    ∂v/∂t = -ω²u - γv + F(u,v) + D₂∇²v
"""
struct CoupledOscillatorPDE <: AbstractPDE
    omega::Float64           # Natural frequency
    gamma::Float64           # Damping coefficient  
    D1::Float64             # Spatial coupling for u
    D2::Float64             # Spatial coupling for v
    forcing::Function       # External forcing F(u,v)
    bc::String
    
    function CoupledOscillatorPDE(omega, gamma=0.1, D1=0.1, D2=0.1, 
                                 forcing=(u,v)->(0.0,0.0); bc="auto_periodic_neumann")
        new(omega, gamma, D1, D2, forcing, bc)
    end
end

"""
    evolution_rate(pde::CoupledOscillatorPDE, state::FieldCollection, t::Float64)

Evolution rate for coupled oscillator system.
"""
function evolution_rate(pde::CoupledOscillatorPDE, state::FieldCollection, t::Float64)
    if length(state) != 2
        throw(ArgumentError("CoupledOscillatorPDE requires exactly 2 fields: [u, v]"))
    end
    
    u, v = state[1], state[2]
    
    # Spatial coupling terms
    u_spatial = pde.D1 * laplacian(u)
    v_spatial = pde.D2 * laplacian(v)
    
    # Oscillator dynamics
    u_data, v_data = data(u), data(v)
    
    # External forcing
    F1, F2 = pde.forcing(u_data, v_data)
    
    # System equations
    dudt = ScalarField(v_data, grid(u)) + u_spatial
    dvdt = ScalarField(-pde.omega^2 * u_data - pde.gamma * v_data + F1, grid(v)) + v_spatial
    
    return FieldCollection([dudt, dvdt])
end

"""
    get_total_concentration(state::FieldCollection)

Compute total concentration for conservation analysis.
"""
function get_total_concentration(state::FieldCollection)
    total = 0.0
    for field in state
        total += sum(data(field))
    end
    return total
end

"""
    get_field_energies(state::FieldCollection)

Compute individual field energies: E_i = ∫|u_i|² dV
"""
function get_field_energies(state::FieldCollection)
    energies = Float64[]
    for field in state
        energy = sum(data(field).^2) / 2
        push!(energies, energy)
    end
    return energies
end

"""
    get_interaction_energy(state::FieldCollection, interaction_func::Function)

Compute interaction energy between fields.
"""
function get_interaction_energy(state::FieldCollection, interaction_func::Function)
    field_data = [data(field) for field in state]
    return interaction_func(field_data...)
end

# Export multi-field PDE framework
export MultiFieldPDE, ReactionDiffusionPDE, CoupledOscillatorPDE
export FitzHughNagumoReaction, SchnakenbergReaction, GrayScottReaction, BrusselatorReaction
export get_total_concentration, get_field_energies, get_interaction_energy